package com.ics.atable.chat.config;

import com.alibaba.cloud.ai.graph.*;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.graph.state.strategy.ReplaceStrategy;
import com.ics.atable.chat.cache.HotKeyCache;
import com.ics.atable.chat.cache.MultiLevelCache;
import com.ics.atable.chat.common.milvus.MilvusMultiVectorStore;
import com.ics.atable.chat.dispatcher.IntentDispatcher;
import com.ics.atable.chat.node.*;
import com.ics.atable.chat.serializer.ItemQueryStateSerializer;
import com.ics.atable.chat.service.EmbeddingService;
import com.ics.atable.chat.service.EmbeddingsPersistService;
import com.ics.atable.chat.service.SemanticModelService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.HashMap;
import java.util.Map;

import static com.alibaba.cloud.ai.graph.StateGraph.END;
import static com.alibaba.cloud.ai.graph.StateGraph.START;
import static com.alibaba.cloud.ai.graph.action.AsyncEdgeAction.edge_async;
import static com.alibaba.cloud.ai.graph.action.AsyncNodeAction.node_async;
import static com.ics.atable.chat.constant.ItemQueryConstant.*;

/**
 * 指标查询工作流
 *
 * @author haohongbin
 */
@Slf4j
@Configuration
public class ItemQueryConfiguration {

    @Autowired
    private ChatClientServiceFactory chatClientServiceFactory;

    @Autowired
    private  EmbeddingService embeddingService;

    @Autowired
    private MilvusMultiVectorStore milvusMultiVectorStore;

    @Autowired
    private EmbeddingsPersistService embeddingsPersistService;

    @Autowired
    private MultiLevelCache multiLevelCache;

    @Autowired
    private HotKeyCache hotKeyCache;

    @Autowired
    private SemanticModelService semanticModelService;


    @Bean
    public StateGraph itemQueryGraph() throws GraphStateException {
        KeyStrategyFactory keyStrategyFactory = () -> {
            HashMap<String, KeyStrategy> keyStrategyHashMap = new HashMap<>();

            // 条件边控制：跳转下一个节点
            keyStrategyHashMap.put(INTENT_NEXT_NODE, new ReplaceStrategy());
            keyStrategyHashMap.put(STRUCT_TOKEN_NEXT_NODE, new ReplaceStrategy());
            keyStrategyHashMap.put(HUMAN_FEEDBACK_NEXT_NODE, new ReplaceStrategy());
            keyStrategyHashMap.put(SYNONYMS_NEXT_NODE, new ReplaceStrategy());
            keyStrategyHashMap.put(ITEM_RETRIEVAL_NEXT_NODE, new ReplaceStrategy());

            // 用户输入
            keyStrategyHashMap.put(QUERY, new ReplaceStrategy());
            keyStrategyHashMap.put(USER_ID, new ReplaceStrategy());
            keyStrategyHashMap.put(SESSION_ID, new ReplaceStrategy());

            keyStrategyHashMap.put(FEED_BACK_COUNT, new ReplaceStrategy());
            keyStrategyHashMap.put(FEED_BACK, new ReplaceStrategy());

            // 节点输出
            keyStrategyHashMap.put(INTENT_RESULT, new ReplaceStrategy());
            keyStrategyHashMap.put(STRUCT_TOKEN_RESULT, new ReplaceStrategy());
            keyStrategyHashMap.put(ITEM_RETRIEVAL_RESULT, new ReplaceStrategy());

            return keyStrategyHashMap;
        };

        // 创建状态图
        StateGraph stateGraph = new StateGraph(ITEM_QUERY_GRAPH, keyStrategyFactory, new ItemQueryStateSerializer(OverAllState::new))
                .addNode(INIT, node_async(new InitNode()))
                // 添加意图识别结点
                .addNode(INTENT, node_async(new IntentNode(chatClientServiceFactory)))
                // 添加结构化分词结点
                .addNode(STRUCT_TOKEN, node_async(new StructTokenNode(chatClientServiceFactory)))
                .addNode(HUMAN_FEEDBACK, node_async(new HumanFeedbackNode()))
                .addNode(SYNONYMS, node_async(new SynonymsNode(chatClientServiceFactory, hotKeyCache, semanticModelService)))
                // 添加并行化检索结点
                .addNode(ITEM_RETRIEVAL, node_async(new ItemRetrievalNode(embeddingService, embeddingsPersistService,
                        milvusMultiVectorStore, multiLevelCache)));

        // 创建图之间的边
        // 建立开始到意图识别结点的边
        stateGraph.addEdge(START, INIT)
                // 初始化后在给意图识别结点
                .addEdge(INIT, INTENT)
                .addConditionalEdges(INTENT, edge_async(new IntentDispatcher()),
                        Map.of(STRUCT_TOKEN, STRUCT_TOKEN, HUMAN_FEEDBACK, HUMAN_FEEDBACK, END, END))
                // 人类反馈后给意识识别结点
                .addEdge(HUMAN_FEEDBACK, INTENT)
                // 使用条件边控制人类反馈次数
//                .addConditionalEdges(HUMAN_FEEDBACK, edge_async(new HumanFeedbackDispatcher()), Map.of(INTENT, INTENT, END, END))
                // 结构化分词到并行化检索的边
                .addEdge(STRUCT_TOKEN, SYNONYMS)
                .addEdge(SYNONYMS, ITEM_RETRIEVAL)
                .addEdge(ITEM_RETRIEVAL, END);

        GraphRepresentation graphRepresentation = stateGraph.getGraph(GraphRepresentation.Type.PLANTUML,
                "item query workflow graph");

        log.info("\n\n");
        log.info(graphRepresentation.content());
        log.info("\n\n");

        return stateGraph;

    }
}
